{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Opion():\n",
    "    \n",
    "    def __init__(self):\n",
    "            \n",
    "        self.dataroot= r'I:\\irregular holes\\paris_eval_gt' #image dataroot\n",
    "        self.maskroot= r'I:\\irregular holes\\testing_mask_dataset'#mask dataroot\n",
    "        self.batchSize= 1   # Need to be set to 1\n",
    "        self.fineSize=256 # image size\n",
    "        self.input_nc=3  # input channel size for first stage\n",
    "        self.input_nc_g=6 # input channel size for second stage\n",
    "        self.output_nc=3# output channel size\n",
    "        self.ngf=64 # inner channel\n",
    "        self.ndf=64# inner channel\n",
    "        self.which_model_netD='basic' # patch discriminator\n",
    "        \n",
    "        self.which_model_netF='feature'# feature patch discriminator\n",
    "        self.which_model_netG='unet_csa'# seconde stage network\n",
    "        self.which_model_netP='unet_256'# first stage network\n",
    "        self.triple_weight=1\n",
    "        self.name='CSA_inpainting'\n",
    "        self.n_layers_D='3' # network depth\n",
    "        self.gpu_ids=[0]\n",
    "        self.model='csa_net'\n",
    "        self.checkpoints_dir=r'.\\checkpoints' #\n",
    "        self.norm='instance'\n",
    "        self.fixed_mask=1\n",
    "        self.use_dropout=False\n",
    "        self.init_type='normal'\n",
    "        self.mask_type='center'\n",
    "        self.lambda_A=100\n",
    "        self.threshold=5/16.0\n",
    "        self.stride=1\n",
    "        self.shift_sz=1 # size of feature patch\n",
    "        self.mask_thred=1\n",
    "        self.bottleneck=512\n",
    "        self.gp_lambda=10.0\n",
    "        self.ncritic=5\n",
    "        self.constrain='MSE'\n",
    "        self.strength=1\n",
    "        self.init_gain=0.02\n",
    "        self.cosis=1\n",
    "        self.gan_type='lsgan'\n",
    "        self.gan_weight=0.2\n",
    "        self.overlap=4\n",
    "        self.skip=0\n",
    "        self.display_freq=1000\n",
    "        self.print_freq=50\n",
    "        self.save_latest_freq=5000\n",
    "        self.save_epoch_freq=2\n",
    "        self.continue_train=False\n",
    "        self.epoch_count=1\n",
    "        self.phase='train'\n",
    "        self.which_epoch=''\n",
    "        self.niter=20\n",
    "        self.niter_decay=100\n",
    "        self.beta1=0.5\n",
    "        self.lr=0.0002\n",
    "        self.lr_policy='lambda'\n",
    "        self.lr_decay_iters=50\n",
    "        self.isTrain=True\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "from util.data_load import Data_load\n",
    "from models.models import create_model\n",
    "import torch\n",
    "import os\n",
    "import torchvision\n",
    "from torch.utils import data\n",
    "import torchvision.transforms as transforms\n",
    "opt = Opion()\n",
    "transform_mask = transforms.Compose(\n",
    "    [transforms.Resize((opt.fineSize,opt.fineSize)),\n",
    "     transforms.ToTensor(),\n",
    "    ])\n",
    "transform = transforms.Compose(\n",
    "    [\n",
    "     transforms.Resize((opt.fineSize,opt.fineSize)),\n",
    "     transforms.ToTensor(),\n",
    "     transforms.Normalize(mean=[0.5] * 3, std=[0.5] * 3)])\n",
    "\n",
    "dataset_test = Data_load(opt.dataroot, opt.maskroot, transform, transform_mask)\n",
    "iterator_test = (data.DataLoader(dataset_test, batch_size=opt.batchSize,shuffle=True))\n",
    "print(len(dataset_test))\n",
    "model = create_model(opt)\n",
    "total_steps = 0\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "load_epoch=30\n",
    "model.load(load_epoch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "\n",
    "save_dir = './measure/true'\n",
    "if os.path.exists(save_dir) is False:\n",
    "    os.makedirs(save_dir)\n",
    "\n",
    "epoch=1\n",
    "i=0\n",
    "for image, mask in (iterator_test):\n",
    "    iter_start_time = time.time()\n",
    "    image=image.cuda()\n",
    "    mask=mask.cuda()\n",
    "    mask=mask[0][0]\n",
    "    mask=torch.unsqueeze(mask,0)\n",
    "    mask=torch.unsqueeze(mask,1)\n",
    "    mask=mask.byte()\n",
    "\n",
    "    model.set_input(image,mask)\n",
    "    model.set_gt_latent()\n",
    "    model.test()\n",
    "    real_A,real_B,fake_B=model.get_current_visuals()\n",
    "    pic = (torch.cat([real_A, real_B,fake_B], dim=0) + 1) / 2.0\n",
    "    torchvision.utils.save_image(pic, '%s/Epoch_(%d)_(%dof%d).jpg' % (\n",
    "    save_dir, epoch, total_steps + 1, len(dataset_test)), nrow=1)\n",
    "\n",
    "            \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [conda env:pytorch]",
   "language": "python",
   "name": "conda-env-pytorch-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
